import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches


# 1-order differential equation
def F(x, y): return y-2*x/y

# parsing solution of y


def Solve_Y(x): return np.sqrt(1+2*x)


# initial value
initial_x = 0
initial_y = 1

# sep width
h = 0.1
x_min = 0
x_max = 5
x = np.arange(x_min, x_max, h)


def plot_result_com_with_parsing(x, y_, s=5, m='*', label_1='Parsing_Solution', label_2='YOUR_Solution', color=['black', 'red']):
    """plot_result_com_with_parsing

    Args:
        x ([numpy.ndarray]): [raw x]
        y_ ([numpy.ndarray]): [numerical answer]
        s (int, optional): [size of marker]. Defaults to 5.
        m (str, optional): [data marker]. Defaults to '*'.
        label_1 (str, optional): [label name]. Defaults to 'Parsing_Solution'.
        label_2 (str, optional): [label name]. Defaults to 'YOUR_Solution'.
        color (list, optional): [color of legends]. Defaults to ['black','red'].

    Return:
        None

    Author: Junno

    Date: 2022-04-26
    """
    global Solve_Y, h, x_min, x_max, initial_x, initial_
    plt.figure(figsize=(8, 6))
    plt.scatter(x, Solve_Y(x), s=s, marker=m, c=color[0])
    plt.scatter(x, y_, s=s, marker=m, c=color[1])
    plt.xlabel('x')
    plt.ylabel('y')
    Parsing_Solution = mpatches.Patch(
        color=color[0], label=label_1)
    Euler_polygonal_arc_forward_method = mpatches.Patch(
        color=color[1], label=label_2)
    plt.title(
        label_2+'method in range {}-{} with h={:.1f}'.format(x_min, x_max, h))
    plt.legend(handles=[Parsing_Solution,
                        Euler_polygonal_arc_forward_method], loc='best')
    plt.show()


def Euler_Forward(x, F, plot_result=True):
    """前向欧拉差分法

    Args:
        x ([numpy.ndarray]): [x]
        F ([function]): [1st-order differential function of y]
        plot_result (bool, optional): [show result using plt]. Defaults to True.

    Return:
        Euler_polygonal_arc ([numpy.ndarray]): [results of Euler_forward method]

    Author: Junno

    Date: 2022-04-26
    """
    global Solve_Y, h, x_min, x_max, initial_x, initial_y
    Euler_polygonal_arc = []
    for i in range(len(x)):
        if x[i] == initial_x:
            if i != 0:
                raise ValueError("initial value should be first")
            Euler_polygonal_arc.append(initial_y)
        else:
            Euler_polygonal_arc.append(
                Euler_polygonal_arc[i-1]+h*F(x[i-1], Euler_polygonal_arc[i-1]))
    if plot_result:
        plot_result_com_with_parsing(
            x, Euler_polygonal_arc, label_2='Euler_Forward')
    return Euler_polygonal_arc


def Euler_backward(x, F, eps, N=100, plot_result=True):
    """后向欧拉差分法

    Args:
        x ([numpy.ndarray]): [x]
        F ([function]): [1st-order differential function of y]
        eps ([float]): [gaps between initial iteration value and the true value]
        N (int, optional): [maximum iteration times]. Defaults to 100.
        plot_result (bool, optional): [show result using plt]. Defaults to True.

    Return:
        Euler_polygonal_arc_back ([numpy.ndarray]): [results of Euler_backward method]

    Author: Junno

    Date: 2022-04-26
    """
    global Solve_Y, h, x_min, x_max, initial_x, initial_y

    def iter_back_Euler(x, y_, y, F, h, eps=1e-2, N=100):
        # iteration step
        v0 = y-eps*10
        v1 = y
        while N or abs(v1-v0) > eps:
            v1 = y_+h*F(x, v0)
            v0 = v1
            N -= 1
        return y

    Euler_polygonal_arc_back = []

    for i in range(len(x)):
        if x[i] == initial_x:
            if i != 0:
                raise ValueError("initial value should be first")
            Euler_polygonal_arc_back.append(initial_y)
        else:
            Euler_polygonal_arc_back.append(iter_back_Euler(x[i], Euler_polygonal_arc_back[i-1], Euler_polygonal_arc_back[i-1]+h*F(x[i-1], Euler_polygonal_arc_back[i-1]),
                                                            F, h))

    if plot_result:
        plot_result_com_with_parsing(
            x, Euler_polygonal_arc_back, label_2='Euler_backward')

    return Euler_polygonal_arc_back


def Euler_Improved(x, F, plot_result=True):
    """改进的Euler格式

    Args:
        x ([numpy.ndarray]): [x]
        F ([function]): [1st-order differential function of y]
        plot_result (bool, optional): [show result using plt]. Defaults to True.

    Return:
        Euler_improved ([numpy.ndarray]): [results of Euler_improved method]

    Author: Junno

    Date: 2022-04-26
    """
    global Solve_Y, h, x_min, x_max, initial_x, initial_y
    Euler_improved = []
    for i in range(len(x)):
        if x[i] == initial_x:
            if i != 0:
                raise ValueError("initial value should be first")
            Euler_improved.append(initial_y)
        else:
            yp = Euler_improved[i-1]+h*F(x[i-1], Euler_improved[i-1])
            yc = Euler_improved[i-1]+h*F(x[i], yp)
            Euler_improved.append(0.5*(yp+yc))

    if plot_result:
        plot_result_com_with_parsing(
            x, Euler_improved, label_2='Euler_improved')

    return Euler_improved


def Euler_dual_step(x, F, plot_result=True):
    """Euler两步格式(预测-校正Euler格式)

    Args:
        x ([numpy.ndarray]): [x]
        F ([function]): [1st-order differential function of y]
        plot_result (bool, optional): [show result using plt]. Defaults to True.

    Return:
        Euler_dual_step ([numpy.ndarray]): [results of Euler_dual_step method]

    Author: Junno

    Date: 2022-04-26
    """
    global Solve_Y, h, x_min, x_max, initial_x, initial_y
    Euler_dual_step = []

    # use Euler_improved method to calculate y1
    for i in range(len(x)):
        if x[i] == initial_x:
            if i != 0:
                raise ValueError("initial value should be first")
            Euler_dual_step.append(initial_y)
        else:
            if i == 1:
                yp = Euler_dual_step[i-1]+h*F(x[i-1], Euler_dual_step[i-1])
                yc = Euler_dual_step[i-1]+h*F(x[i], yp)
                Euler_dual_step.append(0.5*(yp+yc))
            else:
                yn_partial = F(x[i-1], Euler_dual_step[i-1])
                yp = Euler_dual_step[i-2]+2*h*yn_partial
                yc = Euler_dual_step[i-1]+h/2*(yn_partial+F(x[i], yp))
                Euler_dual_step.append(yc)

    if plot_result:
        plot_result_com_with_parsing(
            x, Euler_dual_step, label_2='Euler_dual_step')

    return Euler_dual_step


def Euler_dual_step_dual_compensate(x, F, plot_result=True):
    """双改进的Euler两步格式

    Args:
        x ([numpy.ndarray]): [x]
        F ([function]): [1st-order differential function of y]
        plot_result (bool, optional): [show result using plt]. Defaults to True.

    Return:
        Euler_dual_step_dual_compensate ([numpy.ndarray]): [results of Euler_dual_step_dual_compensate method]

    Author: Junno

    Date: 2022-04-26
    """

    global Solve_Y, h, x_min, x_max, initial_x, initial_y

    Euler_dual_step_dual_compensate = []
    predictions = [0 for _ in range(2)]
    corrections = [0 for _ in range(2)]
    for i in range(len(x)):
        if x[i] == initial_x:
            if i != 0:
                raise ValueError("initial value should be first")
            Euler_dual_step_dual_compensate.append(initial_y)
        else:
            if i == 1:
                yp = Euler_dual_step_dual_compensate[i-1]+h * \
                    F(x[i-1], Euler_dual_step_dual_compensate[i-1])
                yc = Euler_dual_step_dual_compensate[i-1]+h*F(x[i], yp)
                Euler_dual_step_dual_compensate.append(0.5*(yp+yc))
            else:
                yn_partial = F(x[i-1], Euler_dual_step_dual_compensate[i-1])
                # predict
                p_ = Euler_dual_step_dual_compensate[i-2]+2*h*yn_partial
                predictions.append(p_)
                # compensate
                m_ = p_+4/5*(corrections[i-1]-predictions[i-1])
                # calculate
                m_partial = F(x[i], m_)
                # correct
                c_ = Euler_dual_step_dual_compensate[i-1] + \
                    h/2*(m_partial+yn_partial)
                corrections.append(c_)
                # compensate
                y_ = c_-1/5*(c_-p_)
                Euler_dual_step_dual_compensate.append(y_)

    if plot_result:
        plot_result_com_with_parsing(
            x, Euler_dual_step_dual_compensate, label_2='Euler_dual_step_dual_compensate')

    return Euler_dual_step_dual_compensate


def Runge_Kutta_4(x, F, plot_result=True):
    """4-th order Runge_Kutta

    Args:
        x ([numpy.ndarray]): [x]
        F ([function]): [1st-order differential function of y]
        plot_result (bool, optional): [show result using plt]. Defaults to True.

    Return:
        Runge_Kutta_four_order ([numpy.ndarray]): [results of Runge_Kutta_four_order method]

    Author: Junno

    Date: 2022-04-26
    """

    global Solve_Y, h, x_min, x_max, initial_x, initial_y
    Runge_Kutta_four_order = []
    for i in range(len(x)):
        if x[i] == initial_x:
            if i != 0:
                raise ValueError("initial value should be first")
            Runge_Kutta_four_order.append(initial_y)
        else:
            K_1 = F(x[i-1], Runge_Kutta_four_order[i-1])
            K_2 = F(x[i-1]+h/2, Runge_Kutta_four_order[i-1]+h/2*K_1)
            K_3 = F(x[i-1]+h/2, Runge_Kutta_four_order[i-1]+h/2*K_2)
            K_4 = F(x[i-1]+h, Runge_Kutta_four_order[i-1]+h*K_3)
            Runge_Kutta_four_order.append(
                Runge_Kutta_four_order[i-1]+h/6*(K_1+2*(K_2+K_3)+K_4))

    if plot_result:
        plot_result_com_with_parsing(
            x, Runge_Kutta_four_order, label_2='Runge_Kutta_4_order')

    return Runge_Kutta_four_order


def single_step_Runge_Kutta_4(x, y, h, F):
    """single step Runge_Kutta_4

    Args:
        x ([float]): [x]
        y ([float]): [y]
        h ([float]): [seperation gap]
        F ([function]): [1st-order differential function of y]

    Returns:
        [float]: [result of single_step RUnge-Kutta-4-th]

    Author: Junno

    Date: 2022-04-26
    """
    # x and y are results of the last step
    K_1 = F(x, y)
    K_2 = F(x+h/2, y+h/2*K_1)
    K_3 = F(x+h/2, y+h/2*K_2)
    K_4 = F(x+h, y+h*K_3)
    return y+h/6*(K_1+2*(K_2+K_3)+K_4)


def Adams_Predict_Correct(x, F, plot_result=True):
    """预测-校正 Adams 格式

    Args:
        x ([numpy.ndarray]): [x]
        F ([function]): [1st-order differential function of y]
        plot_result (bool, optional): [show result using plt]. Defaults to True.

    Return:
        Adams_pc ([numpy.ndarray]): [results of Adams_Predict_Correct method]

    Author: Junno

    Date: 2022-04-26
    """
    global Solve_Y, h, x_min, x_max, initial_x, initial_y

    f_n = []
    Adams_pc = []
    for i in range(len(x)):
        if x[i] == initial_x:
            if i != 0:
                raise ValueError("initial value should be first")
            f_n = [F(initial_x, initial_y)]
            Adams_pc.append(initial_y)
        else:
            if i < 4:
                yn = single_step_Runge_Kutta_4(x[i-1], Adams_pc[i-1], h, F)
                Adams_pc.append(yn)
                f_n.append(F(x[i], yn))
            else:
                p_ = Adams_pc[i-1]+h/24 * \
                    (55*f_n[i-1]-59*f_n[i-2]+37*f_n[i-3]-9*f_n[i-4])
                f_n_ = F(x[i], p_)
                c_ = Adams_pc[i-1]+h/24 * \
                    (9*f_n_+19*f_n[i-1]-5*f_n[i-2]+f_n[i-3])
                Adams_pc.append(c_)
                f_n.append(F(x[i], c_))

    if plot_result:
        plot_result_com_with_parsing(
            x, Adams_pc, label_2='Adams_Predict_Correct')

    return Adams_pc


def Miline_Hamming_Predict_Correct(x, F, plot_result=True):
    """预测-校正 Miline_Hamming 格式

    Args:
        x ([numpy.ndarray]): [x]
        F ([function]): [1st-order differential function of y]
        plot_result (bool, optional): [show result using plt]. Defaults to True.

    Return:
        Miline_Hamming_pc ([numpy.ndarray]): [results of Miline_Hamming_Predict_Correct method]

    Author: Junno

    Date: 2022-04-26
    """
    global Solve_Y, h, x_min, x_max, initial_x, initial_y
    f_n = []
    Miline_Hamming_pc = []
    for i in range(len(x)):
        if x[i] == initial_x:
            f_n = [F(initial_x, initial_y)]
            Miline_Hamming_pc.append(initial_y)
        else:
            if i < 4:
                yn = single_step_Runge_Kutta_4(
                    x[i-1], Miline_Hamming_pc[i-1], h, F)
                Miline_Hamming_pc.append(yn)
                f_n.append(F(x[i], yn))
            else:
                p_ = Miline_Hamming_pc[i-4]+4*h/3 * \
                    (2*f_n[i-1]-f_n[i-2]+2*f_n[i-3])
                f_n_ = F(x[i], p_)
                c_ = 1/8 * \
                    (9*Miline_Hamming_pc[i-1]-Miline_Hamming_pc[i-3]
                     )+3*h/8*(f_n_+2*f_n[i-1]-f_n[i-2])
                Miline_Hamming_pc.append(c_)
                f_n.append(F(x[i], c_))

    if plot_result:
        plot_result_com_with_parsing(
            x, Miline_Hamming_pc, label_2='Miline_Hamming_Predict_Correct')

    return Miline_Hamming_pc


def Miline_Hamming_dual_step_dual_compensate(x, F, plot_result=True):
    """双校正预测-校正 Miline_Hamming 格式

    Args:
        x ([numpy.ndarray]): [x]
        F ([function]): [1st-order differential function of y]
        plot_result (bool, optional): [show result using plt]. Defaults to True.

    Return:
        Miline_Hamming_dual_compensate ([numpy.ndarray]): [results of Miline_Hamming_dual_step_dual_compensate method]

    Author: Junno

    Date: 2022-04-26
    """
    global Solve_Y, h, x_min, x_max, initial_x, initial_y
    Miline_Hamming_dual_compensate = []
    predictions = [0 for _ in range(4)]
    corrections = [0 for _ in range(4)]
    f_n = []
    for i in range(len(x)):
        if x[i] == initial_x:
            Miline_Hamming_dual_compensate.append(initial_y)
            f_n.append(F(initial_x, initial_y))
        else:
            if i < 4:
                yn = single_step_Runge_Kutta_4(
                    x[i-1], Miline_Hamming_dual_compensate[i-1], h, F)
                Miline_Hamming_dual_compensate.append(yn)
                f_n.append(F(x[i], yn))
            else:
                # predict
                p_ = Miline_Hamming_dual_compensate[i-4] + \
                    4*h/3*(2*f_n[i-1]-f_n[i-2]+2*f_n[i-3])
                predictions.append(p_)
                # compensate
                m_ = p_+112/121*(corrections[i-1]-predictions[i-1])
                # calculate
                m_partial = F(x[i], m_)
                # correct
                c_ = 1/8*(9*Miline_Hamming_dual_compensate[i-1]-Miline_Hamming_dual_compensate[i-3])+3*h/8*(
                    m_partial+2*f_n[i-1]-f_n[i-2])
                corrections.append(c_)
                # compensate
                y_ = c_-9/121*(c_-p_)
                Miline_Hamming_dual_compensate.append(y_)
                f_n.append(F(x[i], y_))

    if plot_result:
        plot_result_com_with_parsing(
            x, Miline_Hamming_dual_compensate, label_2='Miline_Hamming_dual_step_dual_compensate')

    return Miline_Hamming_dual_compensate


if __name__ == "__main__":
    """
    ******测试提示******
    Needed python pkgs: numpy、matplotlib
    python: version>3.0
    请按照STEP的提示取消注释相应的代码块
    进行第n步STEP代码测试时请注释掉其他STEP代码，保存修改后执行，并查看结果图像
    执行建议1：可以在cmd或者prompt中跳转到指定目录下，然后输入"python xxx.py"即可执行该python文件，
    前提是测试电脑上已装好相应的python环境以及库
    执行建议2：可以下载vscode或者spyder进行测试，相应步骤请参考度娘!!
    所有代码输出的图例都有标题，如果标题有误可修改源代码中plot部分的'label_2'选项
    **原创声明：以上代码均为个人编写，相应的原创证明已发表在在csdn blogs，请测试教务助理不要二次传播，仅做教学作业检验使用！**
    """

    SHOW_RESULT = True

    # *****STEP 0*****
    # check raw x and y
    # plt.scatter(x,Solve_Y(x))
    # plt.show()

    # *****STEP 1*****
    # test Euler forward
    # Euler_forward_result = Euler_Forward(x, F, SHOW_RESULT)

    # *****STEP 2*****
    # test Euler backward
    # Euler_backward_result=Euler_backward(x, F, 1e-2,SHOW_RESULT)

    # *****STEP 3*****
    # test Euler-Improved
    # Euler_improved_result = Euler_Improved(x, F, SHOW_RESULT)

    # *****STEP 4*****
    # test Euler-Improved
    # Euler_dual_step_result = Euler_dual_step(x, F, SHOW_RESULT)

    # *****STEP 5*****
    # test Euler_dual_step_dual_compensate method
    # Euler_dual_step_dual_compensate_result = Euler_dual_step_dual_compensate(
    #     x, F, SHOW_RESULT)

    # *****STEP 6*****
    # test Runge_Kutta_four_order
    # Runge_Kutta_4_result = Runge_Kutta_4(x, F, SHOW_RESULT)

    # *****STEP 7*****
    # comparison of methods above
    # SHOW_RESULT = False
    # plt.figure(figsize=(8, 6))
    # plt.scatter(x, Solve_Y(x), s=6, marker='*', c='black')
    # plt.scatter(x, Euler_forward_result, s=6, marker='*', c='blue')
    # plt.scatter(x, Euler_improved_result, s=6, marker='*', c='red')
    # plt.scatter(x, Euler_dual_step_result, s=6, marker='*', c='green')
    # plt.scatter(x, Euler_dual_step_dual_compensate_result,
    #             s=6, marker='*', c='orange')
    # plt.scatter(x, Runge_Kutta_4_result, s=6, marker='*', c='pink')
    # plt.xlabel('x')
    # plt.ylabel('y')
    # Parsing_Solution = mpatches.Patch(color='black', label='Parsing_Solution')
    # Euler_forward_match = mpatches.Patch(color='blue', label='Euler_forward')
    # Euler_improved_match = mpatches.Patch(color='red', label='Euler_improved')
    # Euler_dual_step_match = mpatches.Patch(
    #     color='green', label='Euler_dual_step')
    # Euler_dual_step_dual_compensate_match = mpatches.Patch(
    #     color='orange', label='Euler_dual_step_dual_compensate')
    # Runge_Kutta_four_order_match = mpatches.Patch(
    #     color='pink', label='Runge_Kutta_four_order')

    # plt.title('Comparison of different numerical deviarion methods in range {}-{} with h={:.1f}'.format(x_min, x_max, h))
    # plt.legend(handles=[Parsing_Solution, Euler_forward_match, Euler_improved_match, Euler_dual_step_match,
    #            Euler_dual_step_dual_compensate_match, Runge_Kutta_four_order_match], loc='best')
    # plt.show()

    # *****STEP 8*****
    # test Admas-pc
    # Adams_pc = Adams_Predict_Correct(x, F, SHOW_RESULT)

    # *****STEP 9*****
    # test Miline_Hamming-pc
    # Miline_Hamming_pc = Miline_Hamming_Predict_Correct(x, F, SHOW_RESULT)

    # *****STEP 10*****
    # test Miline_Hamming_dual_step_dual_compensate
    # Miline_Hamming_dual_compensate = Miline_Hamming_dual_step_dual_compensate(
    #     x, F, SHOW_RESULT)

    # *****STEP 11*****
    # comparison of 4-th methods above with different h by infinity-error
    # SHOW_RESULT = False
    # seps = np.arange(0.1, 1, 0.05)
    # Runge_Kutta_Infty_errors = []
    # Adams_pc_Infty_errors = []
    # Miline_Hamming_pc_Infty_errors = []
    # Miline_Hamming_pc_dual_compensate_Infty_errors = []
    # for h in seps:
    #     # recreate x with given h
    #     x = np.arange(x_min, x_max, h)
    #     # 4-th order methods
    #     Runge_Kutta_4_result = Runge_Kutta_4(x, F, SHOW_RESULT)
    #     Adams_pc = Adams_Predict_Correct(x, F, SHOW_RESULT)
    #     Miline_Hamming_pc = Miline_Hamming_Predict_Correct(x, F, SHOW_RESULT)
    #     Miline_Hamming_dual_compensate = Miline_Hamming_dual_step_dual_compensate(
    #         x, F, SHOW_RESULT)
    #     # calculate Infty_errors
    #     Runge_Kutta_Infty_errors.append(
    #         np.max(np.abs(Solve_Y(x)-Runge_Kutta_4_result)))
    #     Adams_pc_Infty_errors.append(np.max(np.abs(Solve_Y(x)-Adams_pc)))
    #     Miline_Hamming_pc_Infty_errors.append(
    #         np.max(np.abs(Solve_Y(x)-Miline_Hamming_pc)))
    #     Miline_Hamming_pc_dual_compensate_Infty_errors.append(
    #         np.max(np.abs(Solve_Y(x)-Miline_Hamming_dual_compensate)))

    # # show result
    # plt.figure(figsize=(8, 6))
    # plt.plot(seps, Runge_Kutta_Infty_errors, marker='*', c='red')
    # plt.plot(seps, Adams_pc_Infty_errors, marker='*', c='blue')
    # plt.plot(seps, Miline_Hamming_pc_Infty_errors, marker='*', c='green')
    # plt.plot(seps, Miline_Hamming_pc_dual_compensate_Infty_errors,
    #          marker='*', c='violet')

    # Runge_Kutta_Infty_errors_match = mpatches.Patch(
    #     color='red', label='Runge_Kutta_4_th')
    # Adams_pc_Infty_errors_match = mpatches.Patch(
    #     color='blue', label='Adams_Predict_Correct')
    # Miline_Hamming_pc_Infty_errors_match = mpatches.Patch(
    #     color='green', label='Miline_Hamming_Predict_Correct')
    # Miline_Hamming_pc_dual_compensate_Infty_errors_match = mpatches.Patch(
    #     color='violet', label='Miline_Hamming_dual_compensate')
    # plt.xlabel('h')
    # plt.ylabel('Infinity Error')
    # plt.title('Infinity errors of different 4-th order methods in range {}-{} with different h'.format(x_min, x_max))
    # plt.legend(handles=[Runge_Kutta_Infty_errors_match, Adams_pc_Infty_errors_match,
    #                     Miline_Hamming_pc_Infty_errors_match, Miline_Hamming_pc_dual_compensate_Infty_errors_match], loc='best')
    # plt.show()
